package scripts

import (
	"bytes"
	"encoding/binary"
	"github.com/jweny/pocassist/pkg/util"
	"net"
	"time"
)

// []byte 转int BytesToInt([]byte{0,0,3,232}) //1000
func BytesToInt(b []byte) (r int, err error) {
	bytesBuffer := bytes.NewBuffer(b)
	var x int32
	err = binary.Read(bytesBuffer, binary.BigEndian, &x)
	if err != nil {
		return
	}
	r = int(x)
	return
}

func CheckSmbCve20200796(args *ScriptScanArgs) (*util.ScanResult, error) {
	//socket地址
	addr := args.Host + ":445"
	conn, err := net.DialTimeout("tcp", addr, time.Second*3)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	//向服务端发送数据
	payload := []byte("\x00\x00\x00\xc0\xfeSMB@\x00\x00\x00\x00\x00\x00\x00\x00\x00\x1f\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00$\x00\x08\x00\x01\x00\x00\x00\x7f\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00x\x00\x00\x00\x02\x00\x00\x00\x02\x02\x10\x02\"\x02$\x02\x00\x03\x02\x03\x10\x03\x11\x03\x00\x00\x00\x00\x01\x00&\x00\x00\x00\x00\x00\x01\x00 \x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x03\x00\n\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00")
	conn.SetWriteDeadline(time.Now().Add(time.Second * 5))
	conn.Write(payload)
	//处理返回
	tmp := make([]byte, 4)
	tempData, err := conn.Read(tmp)
	if err != nil {
		return nil, err
	}
	length, err := BytesToInt(tmp[:tempData])
	if err != nil {
		return nil, err
	}

	buffer := make([]byte, length)
	bufferData, err := conn.Read(buffer)
	if err != nil {
		return nil, err
	}
	receData := buffer[:bufferData]
	num1 := bytes.Compare(receData[68:70], []byte("\x11\x03"))
	num2 := bytes.Compare(receData[70:72], []byte("\x02\x00"))
	//num == 0 时为相等
	if num1 != 0 || num2 != 0{
		return &util.InVulnerableResult, nil
	} else {
		return util.VulnerableTcpOrUdpResult(addr,"",nil,nil), nil
	}
}

func init() {
	ScriptRegister("poc-go-smb-cve-2020-0796", CheckSmbCve20200796)
}
